from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from agents.components.base import BaseComponent, ComponentConfig
from agents.components.vector.chroma_component import ChromaComponent
from agents.components.embed.base_embedding_component import EmbeddingComponent

@dataclass
class Connection:
    """组件连接定义"""
    source_component: str
    target_component: str
    source_output: str
    target_input: str

class Flow:
    """流程编排类"""
    
    def __init__(self, name: str):
        self.name = name
        self.components: Dict[str, BaseComponent] = {}
        self.connections: List[Connection] = []
        self.execution_order: List[str] = []
        
    def add_component(self, component_id: str, component: BaseComponent):
        """添加组件"""
        self.components[component_id] = component
        
        # 如果是ChromaComponent，需要设置其embedding组件
        if isinstance(component, ChromaComponent):
            embedding_component_id = component.embedding_component_id
            if embedding_component_id in self.components:
                embedding_component = self.components[embedding_component_id]
                if isinstance(embedding_component, EmbeddingComponent):
                    component.set_embedding_component(embedding_component)
                else:
                    raise ValueError(f"组件 {embedding_component_id} 不是EmbeddingComponent类型")
            else:
                raise ValueError(f"未找到embedding组件 {embedding_component_id}")
        
    def add_connection(self, connection: Connection):
        """添加连接"""
        self.connections.append(connection)
        
    def _build_execution_order(self):
        """构建执行顺序"""
        # 构建邻接表表示的有向图
        graph = {component_id: [] for component_id in self.components.keys()}
        in_degree = {component_id: 0 for component_id in self.components.keys()}
        
        # 根据连接关系构建图
        for conn in self.connections:
            if conn.source_component in self.components and conn.target_component in self.components:
                graph[conn.source_component].append(conn.target_component)
                in_degree[conn.target_component] += 1
        
        # 添加ChromaComponent和其embedding组件的依赖关系
        for component_id, component in self.components.items():
            if isinstance(component, ChromaComponent):
                embedding_component_id = component.embedding_component_id
                if embedding_component_id in self.components:
                    graph[embedding_component_id].append(component_id)
                    in_degree[component_id] += 1
        
        # 拓扑排序
        queue = []
        # 找出所有入度为0的节点
        for component_id, degree in in_degree.items():
            if degree == 0:
                queue.append(component_id)
        
        execution_order = []
        while queue:
            current = queue.pop(0)
            execution_order.append(current)
            
            # 更新相邻节点的入度
            for neighbor in graph[current]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)
        
        # 检查是否存在循环依赖
        if len(execution_order) != len(self.components):
            raise ValueError("流程中存在循环依赖，无法确定执行顺序")
            
        self.execution_order = execution_order
        
    async def execute(self, initial_input: Dict[str, Any]) -> Dict[str, Any]:
        """执行流程"""
        # 构建执行顺序
        self._build_execution_order()

        print("execution_order: ", self.execution_order)
        
        # 初始化上下文
        context = {}  # 不再直接复制initial_input到context
        
        # 按顺序执行组件
        for component_id in self.execution_order:
            component = self.components[component_id]
            
            # 收集组件输入
            component_input = {}
            
            # 1. 首先检查是否有来自initial_input的输入
            for key, value in initial_input.items():
                print(component.get_input_schema())
                if key in component.get_input_schema()['properties']:  # 假设组件有一个方法返回其输入schema
                    component_input[key] = value
            
            # 2. 然后检查是否有来自其他组件的输入
            for conn in self.connections:
                if conn.target_component == component_id:
                    # 从context中获取带前缀的输入
                    prefixed_output = f"{conn.source_component}_{conn.source_output}"
                    component_input[conn.target_input] = context.get(prefixed_output)
            
            print("component_input: ", component_input)
            print("component_id: ", component_id)
            
            # 执行组件
            output = await component.execute(component_input)
            
            print("output: ", output)

            # 更新上下文，为输出添加组件ID前缀
            prefixed_output = {f"{component_id}_{k}": v for k, v in output.items()}
            context.update(prefixed_output)
            
        return context 